import os

import torch
from sklearn.metrics import accuracy_score
from transformers import AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments
from datasets import load_dataset
from safetensors.torch import load_file

# Check if CUDA is available and set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

os.environ["HF_DATASETS_CACHE"] = "D:/Project/AI/cache/huggingface_datasets"

# Load model and tokenizer
model_path = "D:/Project/AI/GoogleBERT/bert-base-multilingual-cased"
tokenizer = AutoTokenizer.from_pretrained(model_path)

#创建一个预训练的模型，并指定它用于 三分类 任务（num_labels=23。这就是模型的类别数，决定了模型输出层的大小（即输出三个类别的 logits）。这也是 CrossEntropyLoss 计算损失时所期望的类别数。也就是n_classes。
model = AutoModelForSequenceClassification.from_pretrained(
    model_path, num_labels=3)

# Move model to GPU if available
model = model.to(device)

# 启用梯度检查点
model.gradient_checkpointing_enable()

# Load and preprocess data

# 使用 XNLI 数据集进行微调，XNLI 数据集是一个 三分类任务，XNLI 数据集本身是一个 跨语言的自然语言推理（NLI）任务，其中包含了 15 种语言，且每个示例都有一个标签，该标签表示句子对之间的关系，包括以下 3 个分类标签：
# 0: entailment（蕴含）
# 1: contradiction（矛盾）
# 2: neutral（中立）
dataset_path = "D:/Project/AI/datasets/facebook/xnli"
dataset = load_dataset(dataset_path, "all_languages")
# 随机抽取部分训练数据（比如10%），注意这里不会覆盖原始的 dataset 字典
train_subset = dataset["train"].shuffle(seed=42).select([i for i in range(1000)])
# 同样的，选择测试集
test_subset = dataset["test"]

# 数据预处理：tokenize
def tokenize_function(examples):
    # 提取英文翻译作为训练文本
    premises = [premise["en"] for premise in examples["premise"]]  # 获取英文的 premise
    hypotheses = [hypothesis["translation"][4] for hypothesis in examples["hypothesis"]] # 获取英文的 hypothesis（索引4是英文）
    
    # 使用 tokenizer 进行批量处理
    # 使用 tokenizer 对 premises 和 hypotheses 进行分词，padding="max_length" 表示将文本填充到最大长度，truncation=True 表示截断超过最大长度的文本。
    # 定义tokenization函数，将 max_length 设置为128
    tokenized_output = tokenizer(
        premises, hypotheses, padding="max_length", truncation=True, max_length=128
    )

    # 获取标签并存储到tokenized_output
    labels = examples["label"]  # 对应每个示例的标签
    tokenized_output["labels"] = labels

    return tokenized_output

# 对数据集的每个子集应用tokenization
tokenized_train_dataset = train_subset.map(tokenize_function, batched=True)
tokenized_test_dataset = test_subset.map(tokenize_function, batched=True)

# 设置格式，使得数据可以被Trainer正确加载
tokenized_train_dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "labels"])
tokenized_test_dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "labels"])

def compute_metrics(p):
    predictions, labels = p
    preds = predictions.argmax(axis=-1)  # 获取预测类别
    return {'eval_accuracy': accuracy_score(labels, preds)}  # 计算并返回 'eval_accuracy'

# Set training arguments
training_args = TrainingArguments(
    output_dir="./results", # 训练结果保存的目录。
    # num_train_epochs=3, # 训练的轮数（epoch）
    num_train_epochs=1,  # 设置较少的训练轮次
    # per_device_train_batch_size=8, 
    # per_device_eval_batch_size=8,  # 每个设备（GPU/CPU）上的批处理大小
    per_device_train_batch_size=4,  # 调小batch size
    per_device_eval_batch_size=4,
    gradient_accumulation_steps=8,  # 每8步进行一次梯度更新，相当于批次大小为 32
    evaluation_strategy="epoch",  # 指定评估策略，此处为每个 epoch 后评估一次。
    logging_dir='./logs',  # 日志文件保存目录。
    logging_steps=100,  # 每 100 步记录日志。
    save_strategy="epoch",  # 每个 epoch 后保存模型。
    load_best_model_at_end=True,  # 训练结束后加载最优模型。
    metric_for_best_model="eval_accuracy",  # 使用 'eval_accuracy'，确保这里和你自定义的指标一致
    # 其他参数
    fp16=True,  # Use 16-bit mixed precision training to speed up training
    dataloader_pin_memory=True,  # 将数据加载到固定内存中，加速数据加载。
    # dataloader_num_workers=0,  # Disable multiprocessing for DataLoader   
)

# Train
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_train_dataset,  # 使用处理后的训练集
    eval_dataset=tokenized_test_dataset,    # 使用处理后的测试集
    compute_metrics=compute_metrics,  # 设置自定义计算指标
)

if __name__ == "__main__":
    trainer.train()
